// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/common/faulty_file_system.h>
#include <melon/synchronization/call_once.h>

#include <filesystem>

namespace kumo::pollux::tests::utils {
namespace {
// Constructs the faulty file path based on the delegated read file 'path'. It
// pre-appends the faulty file system scheme.
inline std::string faultyPath(const std::string& path) {
  return fmt::format("{}{}", FaultyFileSystem::scheme(), path);
}

std::function<bool(std::string_view)> schemeMatcher() {
  // Note: presto behavior is to prefix local paths with 'file:'.
  // Check for that prefix and prune to absolute regular paths as needed.
  return [](std::string_view filePath) {
    return filePath.find(FaultyFileSystem::scheme()) == 0;
  };
}

melon::once_flag faultFilesystemInitOnceFlag;

std::function<std::shared_ptr<
    FileSystem>(std::shared_ptr<const config::ConfigBase>, std::string_view)>
fileSystemGenerator() {
  return [](std::shared_ptr<const config::ConfigBase> properties,
            std::string_view /*unused*/) {
    // One instance of faulty FileSystem is sufficient. Initializes on first
    // access and reuse after that.
    static std::shared_ptr<FileSystem> lfs;
    melon::call_once(faultFilesystemInitOnceFlag, [&properties]() {
      lfs = std::make_shared<FaultyFileSystem>(std::move(properties));
    });
    return lfs;
  };
}
} // namespace

std::unique_ptr<ReadFile> FaultyFileSystem::openFileForRead(
    std::string_view path,
    const FileOptions& options) {
  const std::string delegatedPath = std::string(extractPath(path));
  auto delegatedFile = getFileSystem(delegatedPath, config_)
                           ->openFileForRead(delegatedPath, options);
  return std::make_unique<FaultyReadFile>(
      std::string(path),
      std::move(delegatedFile),
      [&](FaultFileOperation* op) { maybeInjectFileFault(op); },
      executor_);
}

std::unique_ptr<WriteFile> FaultyFileSystem::openFileForWrite(
    std::string_view path,
    const FileOptions& options) {
  const std::string delegatedPath = std::string(extractPath(path));
  auto delegatedFile = getFileSystem(delegatedPath, config_)
                           ->openFileForWrite(delegatedPath, options);
  return std::make_unique<FaultyWriteFile>(
      std::string(path), std::move(delegatedFile), [&](FaultFileOperation* op) {
        maybeInjectFileFault(op);
      });
}

void FaultyFileSystem::remove(std::string_view path) {
  const std::string delegatedPath = std::string(extractPath(path));
  getFileSystem(delegatedPath, config_)->remove(delegatedPath);
}

void FaultyFileSystem::rename(
    std::string_view oldPath,
    std::string_view newPath,
    bool overwrite) {
  const auto delegatedOldPath = extractPath(oldPath);
  const auto delegatedNewPath = extractPath(newPath);
  getFileSystem(delegatedOldPath, config_)
      ->rename(delegatedOldPath, delegatedNewPath, overwrite);
}

bool FaultyFileSystem::exists(std::string_view path) {
  const auto delegatedPath = extractPath(path);
  return getFileSystem(delegatedPath, config_)->exists(delegatedPath);
}

bool FaultyFileSystem::isDirectory(std::string_view path) const {
  const auto delegatedPath = extractPath(path);
  return getFileSystem(delegatedPath, config_)->isDirectory(delegatedPath);
}

std::vector<std::string> FaultyFileSystem::list(std::string_view path) {
  const auto delegatedDirPath = extractPath(path);
  const auto delegatedFiles =
      getFileSystem(delegatedDirPath, config_)->list(delegatedDirPath);
  // NOTE: we shall return the faulty file paths instead of the delegated file
  // paths for list result.
  std::vector<std::string> files;
  files.reserve(delegatedFiles.size());
  for (const auto& delegatedFile : delegatedFiles) {
    files.push_back(faultyPath(delegatedFile));
  }
  return files;
}

void FaultyFileSystem::mkdir(
    std::string_view path,
    const DirectoryOptions& options) {
  const auto delegatedDirPath = extractPath(path);
  auto delegatedFs = getFileSystem(delegatedDirPath, config_);

  auto op =
      FaultFileSystemMkdirOperation(std::string(delegatedDirPath), options);
  maybeInjectFilesystemFault(&op);
  if (!op.delegate) {
    return;
  }

  delegatedFs->mkdir(delegatedDirPath, options);
}

void FaultyFileSystem::rmdir(std::string_view path) {
  const auto delegatedDirPath = extractPath(path);
  getFileSystem(delegatedDirPath, config_)->rmdir(delegatedDirPath);
}

void FaultyFileSystem::setFilesystemInjectionHook(
    FileSystemFaultInjectionHook hook) {
  std::lock_guard<std::mutex> l(mu_);
  fsInjections_ = FileSystemInjections(std::move(hook));
}

void FaultyFileSystem::setFileSystemInjectionError(
    std::exception_ptr exception,
    std::unordered_set<FaultFileSystemOperation::Type> opTypes) {
  std::lock_guard<std::mutex> l(mu_);
  fsInjections_ =
      FileSystemInjections(std::move(exception), std::move(opTypes));
}

void FaultyFileSystem::setFileInjectionHook(
    FileFaultInjectionHook injectionHook) {
  std::lock_guard<std::mutex> l(mu_);
  fileInjections_ = FileInjections(std::move(injectionHook));
}

void FaultyFileSystem::setFileInjectionError(
    std::exception_ptr error,
    std::unordered_set<FaultFileOperation::Type> opTypes) {
  std::lock_guard<std::mutex> l(mu_);
  fileInjections_ = FileInjections(std::move(error), std::move(opTypes));
}

void FaultyFileSystem::setFileInjectionDelay(
    uint64_t delayUs,
    std::unordered_set<FaultFileOperation::Type> opTypes) {
  std::lock_guard<std::mutex> l(mu_);
  fileInjections_ = FileInjections(delayUs, std::move(opTypes));
}

void FaultyFileSystem::clearFileSystemInjections() {
  std::lock_guard<std::mutex> l(mu_);
  fsInjections_.reset();
}

void FaultyFileSystem::clearFileFaultInjections() {
  std::lock_guard<std::mutex> l(mu_);
  fileInjections_.reset();
}

void FaultyFileSystem::maybeInjectFilesystemFault(
    FaultFileSystemOperation* op) {
  {
    std::lock_guard<std::mutex> l(mu_);
    if (!fsInjections_.has_value()) {
      return;
    }

    auto& injections = fsInjections_.value();
    if (injections.filesystemInjectionHook != nullptr) {
      injections.filesystemInjectionHook(op);
      return;
    }

    if (!injections.opTypes.empty() &&
        injections.opTypes.count(op->type) == 0) {
      return;
    }

    if (injections.directoryException != nullptr) {
      std::rethrow_exception(injections.directoryException);
    }
  }
  return;
}

void FaultyFileSystem::maybeInjectFileFault(FaultFileOperation* op) {
  FileInjections injections;
  {
    std::lock_guard<std::mutex> l(mu_);
    if (!fileInjections_.has_value()) {
      return;
    }
    injections = fileInjections_.value();
  }

  if (injections.fileInjectionHook != nullptr) {
    injections.fileInjectionHook(op);
    return;
  }

  if (!injections.opTypes.empty() && injections.opTypes.count(op->type) == 0) {
    return;
  }

  if (injections.fileException != nullptr) {
    std::rethrow_exception(injections.fileException);
  }

  if (injections.fileDelayUs != 0) {
    std::this_thread::sleep_for(
        std::chrono::microseconds(injections.fileDelayUs));
  }
}

void registerFaultyFileSystem() {
  registerFileSystem(schemeMatcher(), fileSystemGenerator());
}

std::shared_ptr<FaultyFileSystem> faultyFileSystem() {
  return std::dynamic_pointer_cast<FaultyFileSystem>(
      getFileSystem(FaultyFileSystem::scheme(), {}));
}
} // namespace kumo::pollux::tests::utils
